# -*- coding: utf-8 -*-
# Author: Sun Jiawei
# E-mail: sunjiawei@tbea.com
# 通过网络爬虫获取煤和天然气的价格数据
#

import re
import time
import random
import warnings

from bs4.element import Tag
import requests
import pandas as pd

from collections import OrderedDict
from tqdm import tqdm
from bs4 import BeautifulSoup

# 天然气
gas_home_page = 'https://www.cngold.org/tianranqi/'
coal_home_page = 'https://www.cngold.org/meitan/'


def main(home_page: str, fossil_type: str, order: int = 1, all_data=None):
    """Main function."""
    if all_data is None:
        all_data = {}
    html = requests.get(home_page)
    soup = BeautifulSoup(html.content)
    if fossil_type == 'gas':
        _all_data, stop = _gas_main(soup, all_data)
        all_data.update(_all_data)
        if not stop:
            pages = soup.find_all('a', text=re.compile(r'^\d$'))
            pages_no = [int(i.text) for i in pages]
            next_page = pages[pages_no.index(order + 1)].attrs['href']
            return main(next_page, 'gas', order + 1, all_data)
        else:
            return all_data
    elif fossil_type == 'coal':
        _all_data, stop = _coal_main(soup, all_data)
        all_data.update(_all_data)
        if not stop:
            pages = soup.find_all('a', text=re.compile(r'^\d$'))
            pages_no = [int(i.text) for i in pages]
            next_page = pages[pages_no.index(order + 1)].attrs['href']
            return main(next_page, 'coal', order + 1, all_data)
        else:
            return all_data
    else:
        raise


def _coal_main(soup: BeautifulSoup, all_data=None):
    """Utility function for ``main``.

    Given the base url ``home_page`` and retrieve coal price data from
    2021-01-01 to now."""
    stop_flag = False
    if all_data is None:
        all_data = {}
    target_tags = soup.find_all('a', text=re.compile(r'参考价格'))
    for tag in tqdm(target_tags):
        url, date = retrieve_data_from_tag(tag)
        date = pd.to_datetime(date)
        if date <= pd.to_datetime('2021-11-10'):
            warnings.warn('2021年11月10日之前的数据格式较为复杂，暂未实现抽取脚本')
            stop_flag = True
            break
        else:
            data = retrieve_price(url, 'coal')
            all_data[date] = data
    return all_data, stop_flag


def _gas_main(soup: BeautifulSoup, all_data=None):
    """Utility function for ``main``.

    Given the base url ``home_page`` and retrieve gas price data from
    2021-01-01 to now."""
    stop_flag = False
    if all_data is None:
        all_data = {}
    target_tags = soup.find_all('a', text=re.compile('今日天然气市场价格查询'))
    for tag in tqdm(target_tags):
        url, date = retrieve_data_from_tag(tag)
        date = pd.to_datetime(date)
        if date <= pd.to_datetime('2021-1-1'):
            stop_flag = True
            break
        else:
            data = retrieve_price(url, 'gas')
            all_data[date] = data
    return all_data, stop_flag


def retrieve_data_from_tag(tag: Tag):
    """Retrieve url and date from given tag."""
    url = tag.attrs['href']
    text = tag.text.strip()
    fmt = re.compile(r'\d{4}年\d{1,2}月\d{1,2}日')
    date = fmt.search(text).group()
    return url, date.replace('年', '-').replace('月', '-').replace('日', '')


def retrieve_price(url: str, fossil_type: str):
    """Retrieve price data from ``url``."""
    time.sleep(random.random())
    html = requests.get(url)
    soup = BeautifulSoup(html.content)
    if fossil_type == 'gas':
        td_tags = soup.find_all('td')
        td_tags = [i for i in td_tags if not i.attrs and not i.findChildren()]
        data = OrderedDict({
            'product-name': [],
            'price-type': [],
            'price': [],
            'price-unit': [],
            'type': [],
            'zone': [],
            'release-time': [],
        })
        for i in range(7):
            times = 0
            order = i + times * 7
            while order < 42:
                data[list(data.keys())[i]].append(td_tags[order].text.strip())
                times += 1
                order = i + times * 7
    elif fossil_type == 'coal':
        th_tags = [i for i in soup.find_all('th') if 'style' in i.attrs]
        if len(th_tags) == 5:  # 最新网页的价格排版
            old = False
        else:  # 老版网页的价格排版
            old = True

        td_tags = soup.find_all('td')
        data = OrderedDict({
            'source': [],
            'variety': [],
            'price': [],
            'time': []
        })
        if old:
            td_tags = [i for i in td_tags if not i.attrs and not i.findChildren()]
            # 由于老版价格数据存在图片，故此处用时间列的值来作为锚点定位一行数据
            start = 0
            step = start
            fmt = re.compile(r'\d{4}-\d{1,2}-\d{1,2}')
            while True:
                try:
                    td = td_tags[step]
                    text = td.text.strip()
                    if fmt.search(text):
                        end = step
                        data['source'].append(td_tags[start].text.strip())
                        data['variety'].append(td_tags[start+1].text.strip())
                        data['price'].append(td_tags[start+2].text.strip())
                        data['time'].append(td_tags[end].text.strip())
                        start = step + 1
                    step += 1
                except IndexError:
                    break
        else:
            td_tags = [i for i in td_tags if not i.findChildren() and i.text]
            # 每5个为一组
            group_no = 0
            while True:
                try:
                    sub_list = td_tags[5*group_no: 5*(group_no+1)]
                    data['source'].append(sub_list[0].text.strip())
                    data['variety'].append(sub_list[1].text.strip())
                    data['price'].append(sub_list[3].text.strip())
                    data['time'].append(sub_list[4].text.strip())
                    group_no += 1
                except IndexError:
                    break
    else:
        print(url)
        raise
    return data


def gas_dict2df(data_dict: dict):
    """Convert a dict object to a df."""
    data_list = []
    for key in data_dict:
        i = 0
        while True:
            try:
                data_list.append([key, data_dict[key]['product-name'][i],
                                  data_dict[key]['price-type'][i], data_dict[key]['price'][i],
                                  data_dict[key]['price-unit'][i], data_dict[key]['type'][i],
                                  data_dict[key]['zone'][i], data_dict[key]['release-time'][i]])
                i += 1
            except IndexError:
                break
    return pd.DataFrame(data_list,
                        columns=['日期', '产品名称', '报价类型', '产品价格', '价格单位', '规格', '地区', '发布时间'])


def coal_dict2df(data_dict: dict):
    """Convert a dict object to a df."""
    data_list = []
    for key in data_dict:
        i = 0
        while True:
            try:
                data_list.append([key, data_dict[key]['source'][i],
                                  data_dict[key]['variety'][i], data_dict[key]['price'][i],
                                  data_dict[key]['time'][i]])
                i += 1
            except IndexError:
                break
    return pd.DataFrame(data_list,
                        columns=['日期', '产地', '品种', '价格', '发布时间'])


if __name__ == '__main__':
    main(coal_home_page, 'coal')
